import os

from tqdm import tqdm
from constants import FAKE_TYPES
from option import parse_args


def frame_count(crops_path, func, template, args):
    videos = os.listdir(crops_path)
    videos.sort()
    t = tqdm(videos) if args.save else videos
    for video in t:
        video_crops_path = os.path.join(crops_path, video)
        video_frame_count = len(os.listdir(video_crops_path))
        # if video_frame_count < 290:
        #     func(template.format(video, video_frame_count))
        if video_frame_count < 270:
            func(template.format(video, video_frame_count))


def main():
    args = parse_args()
    crops, version = args.crops_dir, args.compression_version
    # # crops_dir = 'diffs'
    # if not args.fake_type:
    #     crops_path = os.path.join(args.root_dir, 'original_sequences', 'youtube', version, crops)
    #     save_file = f'abnormal_frame_count_{version}.dat'
    # else:
    #     crops_path = os.path.join(args.root_dir, 'manipulated_sequences', args.fake_type, version, crops)
    #     save_file = f'abnormal_frame_count_{args.fake_type}_{version}.dat'

    fake_types = ['']
    fake_types.extend(FAKE_TYPES[3:4])
    for fake_type in fake_types:
        if fake_type == '':
            crops_path = os.path.join(args.root_dir, 'original_sequences', 'youtube', version, crops)
            save_file = f'abnormal_frame_count_{version}.dat'
        else:
            crops_path = os.path.join(args.root_dir, 'manipulated_sequences', fake_type, version, crops)
            save_file = f'abnormal_frame_count_{fake_type}_{version}.dat'
        # print(crops_path)
        args.save = True
        if args.save:
            with open(save_file, 'w') as f:
                frame_count(crops_path, f.write, '{},{}\n', args)
        else:
            frame_count(crops_path, print, '{},{}', args)


if __name__ == '__main__':
    # 'Deepfakes', 'Face2Face', 'FaceSwap', 'NeuralTextures', 'FaceShifter'
    main()
